#!/usr/bin/env python
# -*- coding:utf-8 -*-
# @FileName  :文件打包.py
# @Time      :2024/3/13 
# @Author    :CL
# @email     :1037654919@qq.com

import os
import shutil
import zipfile


def get_folder_size(folder_path):
    total_size = 0
    for dirpath, dirnames, filenames in os.walk(folder_path):
        for filename in filenames:
            filepath = os.path.join(dirpath, filename)
            total_size += os.path.getsize(filepath)
    return total_size


# 分批次打包
def batch_zip_folders(input_folder, output_folder, max_size_per_zip):
    current_batch = 1 + len(os.listdir(output_folder))
    current_size = 0
    current_files = []
    folder_sizes = []
    for folder in os.listdir(input_folder):
        folder_path = os.path.join(input_folder, folder)
        if os.path.isdir(folder_path):
            total_size = sum(os.path.getsize(os.path.join(folder_path, f)) for f in os.listdir(folder_path) if
                             os.path.isfile(os.path.join(folder_path, f)))
            folder_sizes.append((folder, total_size))

    # 按文件夹大小排序
    sorted_folders = sorted(folder_sizes, key=lambda x: -x[1])

    # 输出排序后的文件夹大小
    for folder, folder_size in sorted_folders:
        print(f"文件夹 {folder} 的大小为: {folder_size / (1024 * 1024):.2f} MB")

        folder_path = os.path.join(input_folder, folder)
        # folder_size = get_folder_size(folder_path)

        if current_size + folder_size > max_size_per_zip:
            print(f"{name}_{current_batch}.zip is ready")
            print(current_files)
            current_zip_path = os.path.join(output_folder, f'{name}_{current_batch}.zip')
            try:
                with zipfile.ZipFile(current_zip_path, 'w') as zipf:
                    for folder_path in current_files:
                        for root, _, files in os.walk(folder_path):
                            for file in files:
                                file_path = os.path.join(root, file)
                                zipf.write(file_path, arcname=os.path.relpath(file_path, input_folder))
                # 删除原始文件
                for f in current_files:
                    shutil.rmtree(f)
            except:
                pass
            current_files = []
            current_size = 0
            current_batch += 1
        else:
            current_files.append(folder_path)
            current_size += folder_size

    if current_files:
        print(current_files)
        current_zip_path = os.path.join(output_folder, f'{name}_{current_batch}.zip')
        with zipfile.ZipFile(current_zip_path, 'w') as zipf:
            for folder_path in current_files:
                for root, _, files in os.walk(folder_path):
                    for file in files:
                        file_path = os.path.join(root, file)
                        zipf.write(file_path, arcname=os.path.relpath(file_path, input_folder))

        for f in current_files:
            shutil.rmtree(f)


# 删除input_folder下的空文件夹
def delete_empty_folders(input_folder):
    for folder in os.listdir(input_folder):
        folder_path = os.path.join(input_folder, folder)

        if os.path.isdir(folder_path) and not os.listdir(folder_path):  # 检查文件夹是否为空
            os.rmdir(folder_path)
            print(f"文件夹 {folder} 已删除")


# 输出input_folder下每个文件夹大小
def get_folder_sizes(input_folder):
    for folder in os.listdir(input_folder):
        folder_path = os.path.join(input_folder, folder)

        if os.path.isdir(folder_path):
            total_size = sum(os.path.getsize(os.path.join(folder_path, f)) for f in os.listdir(folder_path) if
                             os.path.isfile(os.path.join(folder_path, f)))
            print(f"文件夹 {folder} 的大小为: {total_size / 1024 / 1024 / 1024:.2f} G")


if __name__ == '__main__':
    print()
    # 设置输入和输出文件夹路径以及每个包的最大大小
    name = 'huaban'  # hippopx freejpg  pinterest,不适用于花瓣数据的打包
    input_folder = f'/media/chenglei3/77D014CED257D1411/images/{name}/'
    output_folder = f'/media/chenglei3/77D014CED257D1411/images/{name}_zip/'
    os.makedirs(output_folder, exist_ok=True)
    max_size_per_zip = 1e9  # 960M

    # # 输出input_folder下每个文件夹大小,按照G的单位
    # get_folder_sizes(input_folder)
    # 删除空文件夹
    delete_empty_folders(input_folder)
    # 执行分批次打包
    batch_zip_folders(input_folder, output_folder, max_size_per_zip)
